import torch
import torch.nn as nn
import torch.utils.data as Data
import torch.nn.functional as F
import utils
import random
import pickle

class DataLoader():
    def __init__(self, hyper_params):
        self.hyper_params = hyper_params

        self.train_data_p = hyper_params['train_data_p']
        self.data_type = hyper_params['data_type']
        self.rssi_list = None

        self.load_data()
        self.hyper_params['mac_cnt'] = self.rssi_list.mac_cnt
        self.mac_cnt = self.rssi_list.mac_cnt
        self.mac_map = self.rssi_list.mac_map
        self.hyper_params['max_x'] = self.rssi_list.max_x
        self.hyper_params['max_y'] = self.rssi_list.max_y
        self.hyper_params['device_cnt'] = self.rssi_list.device_cnt

    def load_data(self):
        try:
            self.train_data = torch.load(open(f'./net/{self.data_type}_train.pkl', 'rb'))
            self.test_data = torch.load(open(f'./net/{self.data_type}_test.pkl', 'rb'))
            self.rssi_list = pickle.load(open(f'./net/{self.data_type}_rssi_list.pkl', 'rb'))
        except Exception as err:
            print(err)
            self.load_from_file()

    def get_pos_data(self, x, y):
        data_x = torch.zeros(self.device_cnt, dtype=torch.float)
        data_y = torch.LongTensor([x, y])
        data_mask = torch.zeros(self.device_cnt, dtype=torch.bool)
        cnt = torch.zeros(self.device_cnt, dtype=torch.float)

        for wifi in self.rssi_list.get_list_at_pos(x, y):
            if wifi.RSSI > -30.0 or self.rssi_list.mac_map.get(wifi.MAC) is None:
                continue
            data_x[self.rssi_list.mac_map[wifi.MAC]] += (wifi.RSSI - self.rssi_list.rssi_min_value) / 70.0
            cnt[self.rssi_list.mac_map[wifi.MAC]] += 1.0
            data_mask[self.rssi_list.mac_map[wifi.MAC]] = True

        data_x /= (cnt + 1e-5)
        return data_x, data_y, data_mask


    def generate_path_data(self, total_cnt, path_len):
        data_x = []
        data_y = []
        data_mask = []
        dx, dy = [1, 0, -1, 0], [0, -1, 0, 1]
        stay_p = 0.6

        for cnt in range(total_cnt):
            path_data_x = []
            path_data_y = []
            path_data_mask = []

            # Start Point
            cur_x = random.randint(0, self.rssi_list.max_x-1)
            cur_y = random.randint(0, self.rssi_list.max_y-1)

            for step in range(path_len):
                if random.random() > stay_p:
                    while True:
                        i = random.randint(0, 3)
                        if cur_x+dx[i]>=0 and cur_x+dx[i]<self.rssi_list.max_x and cur_y+dy[i]>=0 and cur_y+dy[i]<self.rssi_list.max_y:
                            cur_x += dx[i]
                            cur_y += dy[i]
                            break
                cur_data_x, cur_data_y, cur_mask = self.get_pos_data(cur_x, cur_y)
                path_data_x.append(cur_data_x)
                path_data_y.append(cur_data_y)
                path_data_mask.append(cur_mask)

            data_x.append(torch.stack(path_data_x))
            data_y.append(torch.stack(path_data_y))
            data_mask.append(torch.stack(path_data_mask))

            if cnt % 50 == 0:
                print(f'{cnt}/{total_cnt} finished.')
        
        return torch.stack(data_x), torch.stack(data_y), torch.stack(data_mask)


    def load_from_file(self):
        # Load the data with RSSIList
        self.rssi_list = utils.init_list(self.hyper_params['data_path'])
        train_data_x = []
        train_mask = []
        train_data_y = []
        test_mask = []
        test_data_x = []
        test_data_y = []
        self.mac_map = self.rssi_list.mac_map
        self.mac_cnt = self.rssi_list.mac_cnt
        self.device_cnt = self.rssi_list.device_cnt

        print('Reading data...')

        if self.data_type == 'normal':
            for i in range(self.rssi_list.max_x):
                for j in range(self.rssi_list.max_y):
                    # Single Train Data
                    cur_train_x, cur_train_y, cur_train_mask = self.get_pos_data(i, j)
                    cur_test_x, cur_test_y, cur_test_mask = self.get_pos_data(i, j)

                    train_data_x.append(cur_train_x)
                    train_mask.append(cur_train_mask)
                    train_data_y.append(cur_train_y)

                    test_data_x.append(cur_test_x)
                    test_mask.append(cur_test_mask)
                    test_data_y.append(cur_test_y)
            train_data_x, train_data_y, train_mask = torch.stack(train_data_x), torch.stack(train_data_y), torch.stack(train_mask)
            test_data_x, test_data_y, test_mask = torch.stack(test_data_x), torch.stack(test_data_y), torch.stack(test_mask)
        elif self.data_type == 'gru':
            train_cnt = self.hyper_params['train_cnt']
            test_cnt = self.hyper_params['test_cnt']
            path_len = self.hyper_params['path_len']

            train_data_x, train_data_y, train_mask = self.generate_path_data(train_cnt, path_len)
            test_data_x, test_data_y, test_mask = self.generate_path_data(test_cnt, path_len)
        else:
            raise NotImplementedError()

        self.train_data = (train_data_x, train_data_y, train_mask)
        self.test_data = (test_data_x, test_data_y, test_mask)
        torch.save(self.train_data, f'./net/{self.data_type}_train.pkl')
        torch.save(self.test_data, f'./net/{self.data_type}_test.pkl')
        pickle.dump(self.rssi_list, open(f'./net/{self.data_type}_rssi_list.pkl', 'wb'))

        print('Finished!')

    def train_dataset(self):
        return Data.TensorDataset(*self.train_data)

    def test_dataset(self):
        return Data.TensorDataset(*self.test_data)
